from itertools import combinations

def load_data():
    # 加载数据集
    return [['bread', 'milk'], ['bread', 'diaper', 'beer', 'eggs'], ['milk', 'diaper', 'beer', 'cola'], ['bread', 'milk', 'diaper', 'beer'], ['bread', 'milk', 'diaper', 'cola']]

def generate_candidates(Lk, k):
    # 生成所有可能的 k 个元素的组合
    Ck = set()
    Lk_list = list(Lk)
    for i in range(len(Lk_list)):
        for j in range(i+1, len(Lk_list)):
            L1 = list(Lk_list[i])[:k-2]
            L2 = list(Lk_list[j])[:k-2]
            L1.sort()
            L2.sort()
            if L1 == L2:  # 如果前 k-2 个元素相同，则合并这两个项集
                Ck.add(frozenset(Lk_list[i] | Lk_list[j]))
    return Ck

def apriori(data, min_support=0.5):
    C1 = {}
    for transaction in data:
        for item in transaction:
            if item not in C1:
                C1[item] = 1
            else:
                C1[item] += 1

    num_items = len(data)
    L1 = []
    support_data = {}
    for item in C1:
        support = C1[item] / num_items
        if support >= min_support:
            L1.append(frozenset([item]))
        support_data[frozenset([item])] = support

    L = [L1]
    k = 2
    while len(L[k-2]) > 0:
        Ck = generate_candidates(L[k-2], k)
        Ck_count = {}
        for transaction in data:
            for candidate in Ck:
                if candidate.issubset(transaction):
                    if candidate not in Ck_count:
                        Ck_count[candidate] = 1
                    else:
                        Ck_count[candidate] += 1
        Lk = []
        for candidate in Ck_count:
            support = Ck_count[candidate] / num_items
            if support >= min_support:
                Lk.append(candidate)
            support_data[candidate] = support
        L.append(Lk)
        k += 1
    return L, support_data

# 示例
data = load_data()
L, support_data = apriori(data, min_support=0.5)

# 格式化输出
print("Frequent Itemsets:")
for k, Lk in enumerate(L):
    if Lk:  # 如果 Lk 不为空
        print(f"  Level {k+1}:")
        for itemset in Lk:
            print(f"    {set(itemset)}")

print("\nSupport Data:")
for itemset, support in support_data.items():
    if support >= 0.5:  # 只输出支持度大于或等于 min_support 的项集
        print(f"  {set(itemset)}: {support:.2f}")
